import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.animation as animation

# ---------------------------
# 参数配置区
# ---------------------------
FLOOR_COUNT = 10          # 总楼层数（推荐≤20）
FLOOR_HEIGHT = 3.5        # 每层高度(m)
CELL_SIZE = 0.5           # 元胞尺寸(m)
STAIR_WIDTH = 2           # 楼梯宽度(m)
STAIR_DEPTH = 1           # 楼梯深度(m)
POPULATION = 50           # 初始总人数
SIMULATION_TIME = 500     # 总模拟时间步

# 体力消耗模型参数
A = 0.71  # 距离衰减系数
B = 0.51  # 时间衰减系数
C_DOWN = 1.0  # 下楼疲劳系数
C_UP = 2.7    # 上楼疲劳系数

# ---------------------------
# 核心类定义
# ---------------------------
class Building:
    """建筑结构类"""
    def __init__(self, floor_count, stair_width, stair_depth):
        self.floor_count = floor_count
        self.stair_positions = []
        self.grid = np.zeros((
            floor_count,
            int(FLOOR_HEIGHT // CELL_SIZE),   # 高度方向元胞数
            int(stair_width // CELL_SIZE),    # 楼梯宽度方向元胞数
            int(stair_depth // CELL_SIZE)     # 楼梯深度方向元胞数
        ), dtype=int)
        
        # 初始化楼梯出口
        for floor in range(floor_count - 1):
            sx = stair_width // CELL_SIZE // 2
            sy = stair_depth // CELL_SIZE // 2
            self.grid[floor, 0, sx, sy] = 2  # 出口标记
            self.stair_positions.append( (floor, sx, sy) )

class Agent:
    """疏散代理类"""
    def __init__(self, floor, x, y):
        self.pos = [floor, x, y]  # [楼层, x坐标, y坐标]
        self.speed = np.random.uniform(0.5, 2.0)
        self.distance = 0
        self.time = 0
        self.exhaustion = 0

class EvacuationSimulator:
    """疏散模拟器"""
    def __init__(self, building, population):
        self.building = building
        self.agents = []
        self.time = 0
        
        # 初始化人群
        for _ in range(population):
            while True:
                floor = np.random.randint(0, building.floor_count-1)
                x = np.random.randint(0, building.grid.shape[2])
                y = np.random.randint(0, building.grid.shape[3])
                if building.grid[floor, 0, x, y] == 0:
                    self.agents.append(Agent(floor, x, y))
                    break

    def update_speed(self, agent):
        """更新移动速度"""
        d = agent.distance
        t = max(agent.time, 1e-3)  # 防止除零
        
        c = C_DOWN if agent.pos[0] < self.building.floor_count-1 else C_UP
        exponent = (d**A) / (t**B)
        factor = 1 - 0.01 * c * np.exp(exponent) + 0.01
        
        return max(0.1, min(agent.speed * factor, 2.0))
    
    def find_neighbors(self, agent):
        """查找可行走邻居"""
        neighbors = []
        for dz in (-1, 0, 1):
            for dx in (-1, 0, 1):
                for dy in (-1, 0, 1):
                    if abs(dz)+abs(dx)+abs(dy) != 1:
                        continue  # 仅允许直线移动
                    nz = agent.pos[0] + dz
                    nx = agent.pos[1] + dx
                    ny = agent.pos[2] + dy
                    if (0 <= nz < self.building.floor_count and
                        0 <= nx < self.building.grid.shape[2] and
                        0 <= ny < self.building.grid.shape[3]):
                        if self.building.grid[nz, 0, nx, ny] == 0:
                            neighbors.append( (nz, nx, ny) )
        return neighbors

    def move_agents(self):
        """执行单步移动"""
        new_agents = []
        for agent in self.agents:
            current_floor = agent.pos[0]
            current_pos = agent.pos[1:]
            
            # 检查是否到达楼梯出口
            if self.building.grid[current_floor, 0, *current_pos] == 2:
                if current_floor < self.building.floor_count-1:
                    # 楼梯转移逻辑
                    next_floor = current_floor + 1
                    sx, sy = self.building.stair_positions[next_floor][:2]
                    new_pos = [next_floor, sx, sy]
                else:
                    continue  # 已到达最终出口
            else:
                # 寻找最优移动方向
                neighbors = self.find_neighbors(agent)
                if not neighbors:
                    new_pos = current_pos
                else:
                    # 优先选择靠近楼梯的方向
                    exit_dir = np.array([
                        self.building.stair_positions[current_floor][1],
                        self.building.stair_positions[current_floor][2]
                    ])
                    distances = [np.linalg.norm(np.array(pos[1:]) - exit_dir) 
                                 for pos in neighbors]
                    new_pos = neighbors[np.argmin(distances)]
            
            # 计算移动距离
            dx = new_pos[1] - current_pos[0]
            dy = new_pos[2] - current_pos[1]
            distance = np.hypot(dx, dy)
            
            # 更新状态
            new_agent = Agent(
                pos=new_pos,
                speed=self.update_speed(agent),
                distance=agent.distance + distance,
                time=agent.time + distance / agent.speed,
                exhaustion=agent.exhaustion + 0.01
            )
            new_agents.append(new_agent)
        
        self.agents = new_agents
        self.time += 1

# ---------------------------
# 可视化模块
# ---------------------------
class Visualizer:
    """3D可视化类"""
    def __init__(self, building, agents):
        self.fig = plt.figure(figsize=(12, 8))
        self.ax = self.fig.add_subplot(111, projection='3d')
        self.building = building
        self.agents = agents
        
        # 初始化图形
        self.ax.set_xlim(0, STAIR_WIDTH)
        self.ax.set_ylim(0, STAIR_DEPTH)
        self.ax.set_zlim(0, FLOOR_COUNT * FLOOR_HEIGHT)
        self.ax.set_xlabel('X (m)')
        self.ax.set_ylabel('Y (m)')
        self.ax.set_zlabel('Z (m)')
        self.ax.set_title('3D Evacuation Simulation')
        
        # 绘制楼梯
        for floor, sx, sy in self.building.stair_positions:
            x = sx * CELL_SIZE
            y = sy * CELL_SIZE
            z_start = floor * FLOOR_HEIGHT
            z_end = (floor+1) * FLOOR_HEIGHT
            self.ax.plot([x, x], [y, y], [z_start, z_end], 
                         color='gray', alpha=0.5)
    
    def update(self, frame):
        """动画更新函数"""
        self.ax.clear()
        self.ax.set_xlim(0, STAIR_WIDTH)
        self.ax.set_ylim(0, STAIR_DEPTH)
        self.ax.set_zlim(0, FLOOR_COUNT * FLOOR_HEIGHT)
        
        # 重新绘制楼梯
        for floor, sx, sy in self.building.stair_positions:
            x = sx * CELL_SIZE
            y = sy * CELL_SIZE
            z_start = floor * FLOOR_HEIGHT
            z_end = (floor+1) * FLOOR_HEIGHT
            self.ax.plot([x, x], [y, y], [z_start, z_end], 
                         color='gray', alpha=0.5)
        
        # 绘制代理
        xs, ys, zs = [], [], []
        colors = []
        for agent in self.agents:
            x = agent.pos[1] * CELL_SIZE
            y = agent.pos[2] * CELL_SIZE
            z = agent.pos[0] * FLOOR_HEIGHT
            xs.append(x)
            ys.append(y)
            zs.append(z)
            colors.append('limegreen' if agent.pos[0]==self.building.floor_count-1 
                          else 'royalblue')
        
        self.ax.scatter(xs, ys, zs, c=colors, s=20, depthshade=False)
        return self.ax.collections

# ---------------------------
# 主程序
# ---------------------------
if __name__ == "__main__":
    # 初始化建筑
    building = Building(FLOOR_COUNT, STAIR_WIDTH, STAIR_DEPTH)
    
    # 初始化模拟器
    simulator = EvacuationSimulator(building, POPULATION)
    
    # 初始化可视化
    vis = Visualizer(building, simulator.agents)
    
    # 创建动画
    ani = animation.FuncAnimation(
        vis.fig, vis.update, 
        frames=SIMULATION_TIME, 
        interval=50, 
        blit=True
    )
    
    plt.show()